#pragma once

// 使用 pair 需要包含的头文件
#include <utility>
using std::pair;
using std::make_pair;
#include <assert.h>
#include <vector>
using std::vector;
using std::cout;
using std::endl;

template <class K, class V>
class AVLTreeNode
{
public:
	// 自引用, 左孩子、右孩子、父亲
	AVLTreeNode<K, V>* _left;
	AVLTreeNode<K, V>* _right;
	AVLTreeNode<K, V>* _parent;

	// _bf : balance factor 平衡因子
	// 平衡因子的值为当前节点的 [右子树高度] - [左子树高度]
	int _bf;
	pair<K, V> _kv;

public:
	AVLTreeNode(const pair<K, V>& kv)
		: _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
		, _kv(kv)
		, _bf(0)
	{}
};

template<typename K, typename V>
class AVLTree
{
	typedef AVLTreeNode<K, V> Node;

public:
	bool insert(const pair<K, V>& kv)
	{
		// 1. 搜索二叉树的 [插入操作]
		if (!_root)
		{
			_root = new Node(kv);
			return true;
		}

		Node* parent = nullptr;
		Node* cur = _root;
		while (cur)
		{
			parent = cur;
			if (kv.first > cur->_kv.first)
				cur = cur->_right;
			else if (kv.first < cur->_kv.first)
				cur = cur->_left;
			else
				return false; // 不允许冗余插入
		}

		cur = new Node(kv);
		if (parent->_kv.first < cur->_kv.first)
			parent->_right = cur;
		else
			parent->_left = cur;
		cur->_parent = parent; // 注意要链接 _parent 指针!!!

		// 2. [往上更新] 更新平衡因子, 如有需要则调整高度
		while (parent)
		{
			// 2.1 cur 插在右边, parent 的平衡因子加一, 反之减一
			if (parent->_right == cur)
				parent->_bf += 1;
			else if (parent->_left == cur)
				parent->_bf -= 1;

			// 2.2 更新完平衡因子后, 如有需要则旋转调整
			// 2.2.1 若 |parent.bf| == 0, 则 _root 整个树已平衡, 结束插入 
			// 2.2.2 若 |parent.bf| == 1, 则 parent 继续往上更新, 直到 _root 为止
			// 2.2.3 若 |parent.bf| == 2, 则当前 parent 子树已不平衡, 需要旋转
			if (parent->_bf == 0)
				break;
			else if (parent->_bf == 1 || parent->_bf == -1)
			{
				parent = parent->_parent;
				cur = cur->_parent;
				continue;
			}
			else if (parent->_bf == 2 || parent->_bf == -2)
			{
				// 【关键】旋转处理后 break
				if (parent->_bf == 2 && cur->_bf == 1)
					rotateL(parent);
				else if (parent->_bf == -2 && cur->_bf == -1)
					rotateR(parent);
				else if (parent->_bf == -2 && cur->_bf == 1)
					rotateLR(parent);
				else if (parent->_bf == 2 && cur->_bf == -1)
					rotateRL(parent);
				else
					assert(false);
				break;
			}
			else
			{
				// 这里说明在 insert 之前, 树已经不是 AVL 树, 出大问题啦!
				assert(false);
			}
		}

		return true;
	}

	void inOrder()
	{
		inOrderHelper(_root);
		cout << endl;
	}

	void preOrder()
	{
		preOrderHelper(_root);
		cout << endl;
	}

	// 判断 _root 是否为 AVL 树
	bool isBalanceTree()
	{
		return isBalanceTreeHelper(_root);
	}

private:
	int getTreeHeight(Node* root)
	{
		if (!root)
			return 0;
		int leftH = getTreeHeight(root->_left);
		int rightH = getTreeHeight(root->_right);
		return leftH > rightH ? leftH + 1 : rightH + 1;
	}

	bool isBalanceTreeHelper(Node* root)
	{
		// 1. 平衡因子没问题
		// 2. AVL 树结构没问题
		if (!root)
			return true;

		int leftH = getTreeHeight(root->_left);
		int rightH = getTreeHeight(root->_right);
		int bf = rightH - leftH;

		if (bf != root->_bf || (bf > 1 || bf < -1))
		{
			return false;
		}

		return isBalanceTreeHelper(root->_left) && isBalanceTreeHelper(root->_right);
	}

	// 左单旋
	void rotateL(Node* parent)
	{
		Node* grandparent = parent->_parent;
		Node* subR = parent->_right;
		Node* subRL = subR->_left;

		// 1. subRL 当 parent 的右孩子
		parent->_right = subRL;
		if (subRL)
			subRL->_parent = parent;
		
		// 2. parent 当 subR 的左孩子
		subR->_left = parent;
		parent->_parent = subR;

		// 3. 处理 grandparent
		subR->_parent = grandparent;
		if (grandparent == nullptr)
		{
			_root = subR;
			_root->_parent = nullptr;
		}
		else
			if (grandparent->_left == parent)
				grandparent->_left = subR;
			else
				grandparent->_right = subR;

		// 4. 已平衡, 平衡因子重置为 0
		parent->_bf = subR->_bf = 0;
	}

	// 右单旋
	void rotateR(Node* parent)
	{
		Node* grandparent = parent->_parent;
		Node* subL = parent->_left;
		Node* subLR = subL->_right;

		// 1. subLR 当 parent 的左孩子
		parent->_left = subLR;
		if (subLR)
			subLR->_parent = parent;

		// 2. parent 当 subL 的右孩子
		subL->_right = parent;
		parent->_parent = subL;

		// 3. 处理 grandparent
		subL->_parent = grandparent;
		if (grandparent == nullptr)
		{
			_root = subL;
			_root->_parent = nullptr;
		}
		else
			if (grandparent->_left == parent)
				grandparent->_left = subL;
			else
				grandparent->_right = subL;

		// 4. 重置平衡因子
		parent->_bf = subL->_bf = 0;
	}

	// 先左旋后右旋
	void rotateLR(Node* parent)
	{
		// 1. 根据插入位置的不同, 更新平衡因子
		Node* subL = parent->_left;
		Node* subLR = subL->_right;
		int subLR_bf = subLR->_bf;

		// 2. 旋转
		rotateL(subL);
		rotateR(parent);

		// 1. 根据插入位置的不同, 更新平衡因子
		if (subLR_bf == 0)
		{
			subL->_bf = 0;
			subLR->_bf = 0;
			parent->_bf = 0;
		}
		else if (subLR_bf == 1)
		{
			subL->_bf = -1;
			subLR->_bf = 0;
			parent->_bf = 0;
		}
		else if (subLR_bf == -1)
		{
			subL->_bf = 0;
			subLR->_bf = 0;
			parent->_bf = 1;
		}
		else
		{
			assert(false);
		}
	}

	// 先右旋再左旋
	void rotateRL(Node* parent)
	{
		// 1. 根据插入位置的不同, 更新平衡因子
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		int subRL_bf = subRL->_bf;

		// 2. 先右旋再左旋
		rotateR(subR);
		rotateL(parent);

		// 1. 根据插入位置的不同, 更新平衡因子
		if (subRL_bf == 0)
		{
			subR->_bf = 0;
			subRL->_bf = 0;
			parent = 0;
		}
		else if (subRL_bf == 1)
		{
			subR->_bf = 0;
			subRL->_bf = 0;
			parent->_bf = -1;
		}
		else if (subRL_bf == -1)
		{
			subR->_bf = 1;
			subRL->_bf = 0;
			parent->_bf = 0;
		}
		else
		{
			assert(false);
		}
	}

	void inOrderHelper(Node* root)
	{
		if (root == nullptr)
			return;
		inOrderHelper(root->_left);
		cout << "(" << root->_kv.first << ", " << root->_bf << ")";
		inOrderHelper(root->_right);
	}

	void preOrderHelper(Node* root)
	{
		if (root == nullptr)
			return;
		cout << "(" << root->_kv.first << ", " << root->_bf << ")";
		preOrderHelper(root->_left);
		preOrderHelper(root->_right);
	}

private:
	Node* _root = nullptr;
};

void test_0()
{
	AVLTree<int, int> t;
	int arr[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
	for (int i = 0; i < sizeof(arr) / sizeof(arr[0]); ++i)
	{
		t.insert(make_pair(arr[i], 0));
	}
	
	if (t.isBalanceTree())
	{
		cout << "t is balance tree.\n";
	}
	else
	{
		cout << "t is not AVL.\n";
	}
}

void test_1()
{
	AVLTree<int, int> t;
	int arr[] = { 4, 2, 6, 1, 3, 5, 15, 7, 16, 14 };
	for (int i = 0; i < sizeof(arr) / sizeof(arr[0]); ++i)
	{
		t.insert(make_pair(arr[i], 0));
	}

	cout << "中序遍历: " << endl;
	t.inOrder();

	if (t.isBalanceTree())
	{
		cout << "t is balance tree.\n";
	}
	else
	{
		cout << "t is not AVL.\n";
	}
}

void test_2()
{
	for(int i = 0; i < 5; ++i)
	{
		srand(time(0));
		const size_t N = 100000;
		AVLTree<int, int> t;
		for (size_t i = 0; i < N; i++)
		{
			size_t x = rand();
			t.insert(make_pair(x, 0));
		}

		cout << t.isBalanceTree() << endl;
	}
}